import sqlite3
import os
import array
import cPickle as pickle
from Crypto.PublicKey import ElGamal
from Crypto.Cipher import AES
from Crypto.Hash import SHA256
from Crypto.Hash import HMAC

_public_key_cache = {}
'''cache to reduce db lookups'''
def generate_keys(db):
	"""\
    Check the database at the supplied path for a public and private key pair.
	Create them if they don't exist.
    """  
	global _public_key_cache
	con = sqlite3.connect(db)
	res = con.execute('select pubkey from my_keys')
	keys = res.fetchall()
	if len(keys) == 0:
		print "No private key found -- please wait while one is created"
		#TODO: increase key size to 2048 after testing is complete
		privateKey = ElGamal.generate(1024, os.urandom)
		publicKey = privateKey.publickey()
		_public_key_cache[db] = publicKey
		strPrivateKey = pickle.dumps(privateKey)
		strPublicKey = pickle.dumps(publicKey)
		
		con.execute('insert into my_keys values (?,?)', 
	                (sqlite3.Binary(strPublicKey), sqlite3.Binary(strPrivateKey)))
		con.commit()
		print "New keys added to database"
		
def get_public_key(db):
	"""\
    Get my public key from the database.  If multiple public keys are found, return
	the first one. 
    """      
	global _public_key_cache
	if db in _public_key_cache:
		return _public_key_cache[db]
	con = sqlite3.connect(db)
	res = con.execute('select pubkey from my_keys')
	keys = res.fetchall()
	if len(keys) == 0:
		raise "no public key found"
	firstRow = keys[0]
	strKey = str(firstRow[0])
	_public_key_cache[db] = pickle.loads(strKey)
	return _public_key_cache[db]
	
def random_symmetric_key():
	"""\
    Generate a random 128-bit symmetric key and associated material.  The resulting
	value should be passed into create_encryption_context to get a context that can 
	be used to perform encryption and decryption.
    """  
	keySize = 16
	key = os.urandom(keySize)
	iv = os.urandom(keySize)
	return (key, AES.MODE_CBC, iv)
	
def create_encryption_context(keyTuple):
	"""\
    Creates an encryption (or decryption) context from a tuple containing key material. 
    """  
	return AES.new(keyTuple[0], keyTuple[1], keyTuple[2])

def encrypt_key(plaintextKey, publicKey):
	"""\
    Encrypts a symmetric key using the supplied public key.
    """  
	randomValue = os.urandom(256)
	#print "encrypting key: " + pickle.dumps(plaintextKey)
	encryptedKey = publicKey.encrypt(pickle.dumps(plaintextKey[0]), randomValue)
	return (encryptedKey, plaintextKey[1], plaintextKey[2])

_key_cache = {}
'''cache to reduce db_lookups'''
def decrypt_key(ciphertextKey, publicKey, db):
	global _key_cache
	"""\
    Decrypts a symmetric key using the private key corresponding to the supplied
	public key.  
    """  
	#print "publicKey = ",publicKey
	pickledPubkey = pickle.dumps(publicKey)
	privateKey = None
	if pickledPubkey in _key_cache:
		privateKey = _key_cache[pickledPubkey]
	else:	
		con = sqlite3.connect(db)
		res = con.execute("select pubkey,private_key from my_keys")
		for row in res:
			if str(row[0]) == pickledPubkey:
				privateKey = pickle.loads(str(row[1]))
	if not privateKey:
		raise Exception("private key not found")	
	#print "got private key: " + pickle.dumps(privateKey)
	decryptedKey = privateKey.decrypt(ciphertextKey[0])
	#print "decrypted key: " + pickle.dumps(result)
	return (pickle.loads(decryptedKey), ciphertextKey[1], ciphertextKey[2])			

 
def add_padding(text):
	"""\
    Takes an arbitrary string and returns a properly padded string that can be supplied
	as an argument to the encrypt functions
    """ 
	length = len(text)
	lengthStr = '{0:016X}'.format(length)
	paddedText = lengthStr + text
	length = len(paddedText)
	requiredPadding = 16 - (length % 16)
	for i in range(requiredPadding):
		paddedText = paddedText + ' '
	return paddedText
	
def remove_padding(text):
	"""\
	Takes a padded string and returns the original string, minus the padding.
	""" 
	lengthStr = text[0:16]
	contentStr = text[16:]
	length = int(lengthStr, 16)
	return contentStr[0:length]
	
def encrypt(plaintext, encryptContext):
	"""\
    Encrypts the supplied text using the supplied encryption context.
    """ 
	# input string must be a multiple of 16 bytes
	paddedPlaintext = add_padding(plaintext)
    #print "encrypting plaintext: " + paddedPlaintext	
	return encryptContext.encrypt(paddedPlaintext)
	
def decrypt(ciphertext, encryptContext):
	"""\
    Decrypts the supplied ciphertext using the supplied encryption context.
    """ 
	paddedPlaintext = encryptContext.decrypt(ciphertext)
	#print "got plaintext: " + paddedPlaintext
	return remove_padding(paddedPlaintext)
	
def hmac(text, key):
	hmacModule = HMAC.new(key[0], None, SHA256)
	hmacModule.update(text)
	return hmacModule.digest()
